#include "kernel_operator.h"
#include <type_traits>
//#include "aclnn_kernels/cast.h
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;


class KernelSelectV2IntInt{
    public :
    __aicore__ inline KernelSelectV2IntInt(){}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t size, uint32_t smallCoreDataNum,
                                uint32_t finalSmallTileNum, 
                                uint32_t tileDataNum,
                                uint32_t smallTailDataNum,
                                uint32_t tailBlockNum){
            ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
            this->tileDataNum = tileDataNum; 
    
            this->coreDataNum = smallCoreDataNum; 
            this->tileNum = finalSmallTileNum;    
            this->tailDataNum = smallTailDataNum;  

            conditionGm.SetGlobalBuffer((__gm__ DTYPE_CONDITION *)condition,this->coreDataNum);
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1,this->coreDataNum);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2,this->coreDataNum);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y,this->coreDataNum);

        }

    __aicore__ inline void process(){
        for(uint32_t i=0; i<this->coreDataNum; i++){
            DTYPE_X1 x1 = x1Gm.GetValue(i);
            DTYPE_X2 x2 = x2Gm.GetValue(i);
            conditionGm.GetValue(i) ? yGm.SetValue(i,x1) : yGm.SetValue(i,x2);
        }
    }
    
    private:
    GlobalTensor<DTYPE_CONDITION> conditionGm;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y>  yGm;

    uint32_t smallCoreDataNum;
    uint32_t tileDataNum;
    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tailDataNum;
};

class KernelSelectV2INT{
    public:

    __aicore__ inline KernelSelectV2INT(){}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t size, uint32_t smallCoreDataNum,
                                uint32_t finalSmallTileNum, 
                                uint32_t tileDataNum,
                                uint32_t smallTailDataNum,
                                uint32_t tailBlockNum){
        ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
        this->tileDataNum = tileDataNum; 

        this->coreDataNum = smallCoreDataNum; 
        this->tileNum = finalSmallTileNum;    
        this->tailDataNum = smallTailDataNum;  

        conditionGm.SetGlobalBuffer((__gm__ DTYPE_CONDITION *)condition,this->coreDataNum);
        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1,this->coreDataNum);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2,this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y,this->coreDataNum);

        pipe.InitBuffer(inQueueCondition,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_CONDITION));
        pipe.InitBuffer(inQueueX1,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(inQueueX2,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_X2));
        pipe.InitBuffer(outQueueY,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_Y));  
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(uint8_t));
        pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp4, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp5, this->tileDataNum * sizeof(half));
    }

    __aicore__ inline void process(){
        int32_t loopcount = this->tileNum;
        this->processDataNum = this->tileDataNum;  
        for (int32_t i = 0; i < loopcount; i++)
        {
            if (i == this->tileNum - 1){
                this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

    private:
    __aicore__ inline void CopyIn(int32_t progress){

        LocalTensor<DTYPE_CONDITION> conditionLocal = inQueueCondition.AllocTensor<DTYPE_CONDITION>();
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
        DataCopy(conditionLocal,conditionGm[progress * this->tileDataNum],(this->processDataNum%32==0?this->processDataNum:((this->processDataNum/32+1)*32)));
        DataCopy(x1Local,x1Gm[progress * this->tileDataNum],this->processDataNum);
        DataCopy(x2Local,x2Gm[progress * this->tileDataNum],this->processDataNum);

        inQueueCondition.EnQue(conditionLocal);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }

    __aicore__ inline void Compute(int32_t progress){
          //将conditionLocal重新解释为uint8类型
        LocalTensor<uint8_t> conditionLocal = inQueueCondition.DeQue<DTYPE_CONDITION>().ReinterpretCast<uint8_t>();
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        auto floatCondition = tmp1.Get<half>();
        auto maskLocal = tmp2.Get<uint8_t>(); 
        LocalTensor<half> x1_float = tmp3.Get<half>();
        LocalTensor<half> x2_float = tmp4.Get<half>();
        LocalTensor<half> y_float = tmp5.Get<half>();
        
        Cast(floatCondition,conditionLocal, RoundMode::CAST_NONE,this->processDataNum);
        // Cast(x1_float,x1Local,RoundMode::CAST_NONE,this->processDataNum);
        // Cast(x2_float,x2Local,RoundMode::CAST_NONE,this->processDataNum);


        CompareScalar(maskLocal, floatCondition, static_cast<half>(1),CMPMODE::EQ,((this->processDataNum+255)/256*256));        

        Select(y_float, maskLocal, x1_float, x2_float,SELMODE::VSEL_TENSOR_TENSOR_MODE,this->processDataNum); 

        // Cast(yLocal,y_float,RoundMode::CAST_ROUND,this->processDataNum);

        outQueueY.EnQue(yLocal);

        inQueueCondition.FreeTensor(conditionLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }

    __aicore__ inline void CopyOut(int32_t progress){
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();

        DataCopy(yGm[progress * this->tileDataNum],yLocal,this->processDataNum);

        outQueueY.FreeTensor(yLocal);
    }



    private:
        TPipe pipe;
        TQue<QuePosition::VECIN,BUFFER_NUM> inQueueCondition,inQueueX1,inQueueX2;
        TQue<QuePosition::VECOUT,BUFFER_NUM> outQueueY;
        TBuf<TPosition::VECCALC> tmp1,tmp2,tmp3,tmp4,tmp5;

        GlobalTensor<DTYPE_CONDITION> conditionGm;
        GlobalTensor<DTYPE_X1> x1Gm;
        GlobalTensor<DTYPE_X2> x2Gm;
        GlobalTensor<DTYPE_Y>  yGm;

        uint32_t smallCoreDataNum;
        uint32_t tileDataNum;
        uint32_t coreDataNum;
        uint32_t tileNum;
        uint32_t tailDataNum;
        uint32_t processDataNum;
};

class KernelSelectV2BC{
    public :
    __aicore__ inline KernelSelectV2BC(){}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t size, uint32_t smallCoreDataNum,
                                uint32_t finalSmallTileNum, 
                                uint32_t tileDataNum,
                                uint32_t smallTailDataNum,
                                uint32_t tailBlockNum, BroadCastTiling BCT){
            ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
            this->tileDataNum = tileDataNum; 
    
            this->coreDataNum = smallCoreDataNum; 
            this->tileNum = finalSmallTileNum;    
            this->tailDataNum = smallTailDataNum;  

            this->dimNum = BCT.y_dimNum;
            
            this->y_dims = BCT.y_dims;
            this->condition_dims = BCT.condition_dims;
            this->x1_dims = BCT.x1_dims;
            this->x2_dims = BCT.x2_dims;

            this->y_sumDims = BCT.y_sumDims;
            this->condition_sumDims = BCT.condition_sumDims;
            this->x1_sumDims = BCT.x1_sumDims;
            this->x2_sumDims = BCT.x2_sumDims;

            conditionGm.SetGlobalBuffer((__gm__ DTYPE_CONDITION *)condition,this->coreDataNum);
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1,this->coreDataNum);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2,this->coreDataNum);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y,this->coreDataNum);

        }

    __aicore__ inline void process(){

        //广播本质：将大小为1的维度复制拓展。
        for(uint32_t i=0; i<this->coreDataNum; i++){
            uint32_t x1_pos = 0, x2_pos = 0, condition_pos = 0;
            for(uint32_t j = 0;j<this->dimNum;j++){
                if(this->condition_dims[j] != 1){
                    condition_pos+=this->condition_sumDims[j] * (i/this->y_sumDims[j]%this->y_dims[j]);
                }
                if(this->x1_dims[j]!=1){
                    x1_pos+=this->x1_sumDims[j] * (i/this->y_sumDims[j]%this->y_dims[j]);
                }
                if(this->x2_dims[j]!=1){
                    x2_pos+=this->x2_sumDims[j] * (i/this->y_sumDims[j]%this->y_dims[j]);
                }
            }
            DTYPE_X1 x1 = x1Gm.GetValue(x1_pos);
            DTYPE_X2 x2 = x2Gm.GetValue(x2_pos);
            conditionGm.GetValue(condition_pos) ? yGm.SetValue(i,x1) : yGm.SetValue(i,x2);
        }
    }
    
    private:
    GlobalTensor<DTYPE_CONDITION> conditionGm;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y>  yGm;

    uint32_t smallCoreDataNum;
    uint32_t tileDataNum;
    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tailDataNum;

    int32_t dimNum;
    int32_t *y_dims, *y_sumDims;
    int32_t *condition_dims, *condition_sumDims;
    int32_t *x1_dims, *x1_sumDims;
    int32_t *x2_dims, *x2_sumDims;
};

class KernelSelectV2 {
    
    public:

    __aicore__ inline KernelSelectV2(){}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t size, uint32_t smallCoreDataNum,
                                uint32_t finalSmallTileNum, 
                                uint32_t tileDataNum,
                                uint32_t smallTailDataNum,
                                uint32_t tailBlockNum){
        ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
        this->tileDataNum = tileDataNum; 

        this->coreDataNum = smallCoreDataNum; 
        this->tileNum = finalSmallTileNum;    
        this->tailDataNum = smallTailDataNum;  

        conditionGm.SetGlobalBuffer((__gm__ DTYPE_CONDITION *)condition,this->coreDataNum);
        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1,this->coreDataNum);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2,this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y,this->coreDataNum);

        pipe.InitBuffer(inQueueCondition,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_CONDITION));
        pipe.InitBuffer(inQueueX1,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(inQueueX2,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_X2));
        pipe.InitBuffer(outQueueY,BUFFER_NUM,this->tileDataNum * sizeof(DTYPE_Y));  
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(uint8_t));
    }

    __aicore__ inline void process(){
        int32_t loopcount = this->tileNum;
        this->processDataNum = this->tileDataNum;  
        for (int32_t i = 0; i < loopcount; i++)
        {
            if (i == this->tileNum - 1){
                this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

    private:
    __aicore__ inline void CopyIn(int32_t progress){

        LocalTensor<DTYPE_CONDITION> conditionLocal = inQueueCondition.AllocTensor<DTYPE_CONDITION>();
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
        DataCopy(conditionLocal,conditionGm[progress * this->tileDataNum],(this->processDataNum%32==0?this->processDataNum:((this->processDataNum/32+1)*32)));
        DataCopy(x1Local,x1Gm[progress * this->tileDataNum],this->processDataNum);
        DataCopy(x2Local,x2Gm[progress * this->tileDataNum],this->processDataNum);

        inQueueCondition.EnQue(conditionLocal);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }

    __aicore__ inline void Compute(int32_t progress){
          //将conditionLocal重新解释为uint8类型
        LocalTensor<uint8_t> conditionLocal = inQueueCondition.DeQue<DTYPE_CONDITION>().ReinterpretCast<uint8_t>();
        
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        auto floatCondition = tmp1.Get<half>();
        auto maskLocal = tmp2.Get<uint8_t>(); 

        Cast(floatCondition,conditionLocal, RoundMode::CAST_NONE,this->processDataNum);
        
        CompareScalar(maskLocal, floatCondition, static_cast<half>(1),CMPMODE::EQ,((this->processDataNum+255)/256*256));        

        Select(yLocal, maskLocal, x1Local, x2Local,SELMODE::VSEL_TENSOR_TENSOR_MODE,this->processDataNum); 

        outQueueY.EnQue(yLocal);
        inQueueCondition.FreeTensor(conditionLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }

    __aicore__ inline void CopyOut(int32_t progress){
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        
        DataCopy(yGm[progress * this->tileDataNum],yLocal,this->processDataNum);

        outQueueY.FreeTensor(yLocal);
    }



    private:
        TPipe pipe;
        TQue<QuePosition::VECIN,BUFFER_NUM> inQueueCondition,inQueueX1,inQueueX2;
        TQue<QuePosition::VECOUT,BUFFER_NUM> outQueueY;
        TBuf<TPosition::VECCALC> tmp1,tmp2;

        GlobalTensor<DTYPE_CONDITION> conditionGm;
        GlobalTensor<DTYPE_X1> x1Gm;
        GlobalTensor<DTYPE_X2> x2Gm;
        GlobalTensor<DTYPE_Y>  yGm;

        uint32_t smallCoreDataNum;
        uint32_t tileDataNum;
        uint32_t coreDataNum;
        uint32_t tileNum;
        uint32_t tailDataNum;
        uint32_t processDataNum;

};

extern "C" __global__ __aicore__ void select_v2(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    if(tiling_data.tilingKey==2){
        KernelSelectV2 op;
        op.Init(condition, x1, x2, y, tiling_data.size,
                                    tiling_data.smallCoreDataNum,
                                    tiling_data.finalSmallTileNum,
                                    tiling_data.tileDataNum,
                                    tiling_data.smallTailDataNum,
                                    tiling_data.tailBlockNum);
        op.process();
    }else if(tiling_data.tilingKey==1){
        KernelSelectV2BC op;
        op.Init(condition, x1, x2, y, tiling_data.size,
                                    tiling_data.smallCoreDataNum,
                                    tiling_data.finalSmallTileNum,
                                    tiling_data.tileDataNum,
                                    tiling_data.smallTailDataNum,
                                    tiling_data.tailBlockNum,
                                    tiling_data.BCT);
        op.process();
    }else if(tiling_data.tilingKey==3){
        KernelSelectV2IntInt op;
        op.Init(condition, x1, x2, y, tiling_data.size,
                                    tiling_data.smallCoreDataNum,
                                    tiling_data.finalSmallTileNum,
                                    tiling_data.tileDataNum,
                                    tiling_data.smallTailDataNum,
                                    tiling_data.tailBlockNum);
        op.process();
    }else{
        KernelSelectV2INT op;
        op.Init(condition, x1, x2, y, tiling_data.size,
                                    tiling_data.smallCoreDataNum,
                                    tiling_data.finalSmallTileNum,
                                    tiling_data.tileDataNum,
                                    tiling_data.smallTailDataNum,
                                    tiling_data.tailBlockNum);
        op.process();

    }
}


    // const int32_t vecLen = (sizeof(DTYPE_X1) == 1) ? 256 :  
    //                           ((sizeof(DTYPE_X1) == 2) ? 128 :  
    //                           64);                             
    
    // const int32_t loopTimes = this->processDataNum / vecLen;  //有多少条指令
    // const int32_t tailElements = this->processDataNum % vecLen;
    
    // for (int32_t i = 0; i < loopTimes; ++i) {
    //     #pragma unroll
    //     for (int32_t j = 0; j < vecLen; ++j) {
    //         const int32_t offset = i * vecLen + j;
    //         printf("x1Local:%f\n",x1Local.GetValue(offset));
    //         printf("x2Local:%f\n",x2Local.GetValue(offset));
    //         printf("conditionLocal:%f\n",conditionLocal.GetValue(offset));
    //         yLocal.SetValue(offset, conditionLocal.GetValue(i*256+j) ? 
    //                       x1Local.GetValue(offset) : x2Local.GetValue(offset));
    //         printf("yLocal:%f\n",yLocal.GetValue(offset));
    //     }
    // }

    // if (tailElements > 0) {
    //     for (int32_t j = 0; j < tailElements; ++j) {
    //         const int32_t offset = loopTimes * vecLen + j;
    //         yLocal.SetValue(offset, conditionLocal.GetValue(loopTimes*256+j) ?
    //                         x1Local.GetValue(offset) : x2Local.GetValue(offset));
    //         printf("！！！！！！！尾块的yLocal:%f\n",yLocal.GetValue(offset));
    //     }
    // }
    
        // for (int32_t i = 0; i < this->processDataNum; ++i) {
        //     DTYPE_CONDITION cond = conditionLocal.GetValue(i);
        //     auto val = cond ? x1Local.GetValue(i) : x2Local.GetValue(i);
        //     yLocal.SetValue(i, val);

        // }
        //uint64_t mask = 256/sizeof(DTYPEX1);
        // // 固定写法，创建OpExecutor
        // auto intCondition = tmp.Get<uint8_t>();
        // uint32_t mask = 256;
        // Copy(intCondition,conditionLocal,mask,
        //     this->processDataNum/mask==0?this->processDataNum/mask:this->processDataNum/mask+1,
        //     {1,1,8,8});
        // for(uint32_t i = 0;i<conditionLocal.GetSize();i++){
        //     // intCondition.SetValue(i, conditionLocal.GetValue(i) == true?1:0);


        // }
        //
